import os
import whisper

VIDEO_ROOT = 'vdeo'
DOC_ROOT = 'doc'


def get_all_mp4_files(root):
    for dirpath, _, filenames in os.walk(root):
        for filename in filenames:
            if filename.lower().endswith('.mp4'):
                yield os.path.join(dirpath, filename)

def get_doc_txt_path(mp4_path):
    rel_path = os.path.relpath(mp4_path, VIDEO_ROOT)
    txt_path = os.path.splitext(rel_path)[0] + '.txt'
    return os.path.join(DOC_ROOT, txt_path)

def ensure_dir_exists(path):
    dir_path = os.path.dirname(path)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path, exist_ok=True)

def transcribe_to_tmp(mp4_path, tmp_txt_path, model):
    # 使用whisper进行转录，增加verbose参数
    result = model.transcribe(mp4_path, language='zh', verbose=True)
    # 将转录文本写入tmp文件
    with open(tmp_txt_path, 'w', encoding='utf-8') as f:
        f.write(result['text'])

def main():
    model = whisper.load_model('medium')
    mp4_files = list(get_all_mp4_files(VIDEO_ROOT))
    total = len(mp4_files)
    for idx, mp4_path in enumerate(mp4_files, 1):
        txt_path = get_doc_txt_path(mp4_path)
        tmp_txt_path = txt_path.replace('.txt', '.tmp.txt')
        if os.path.exists(txt_path):
            print(f"Skip {txt_path} (already exists)")
            continue
        ensure_dir_exists(txt_path)
        print(f"[{idx}/{total}] Transcribing {mp4_path} -> {tmp_txt_path}")
        try:
            transcribe_to_tmp(mp4_path, tmp_txt_path, model)
            os.rename(tmp_txt_path, txt_path)
            print(f"Saved {txt_path}")
        except Exception as e:
            print(f"Failed to transcribe {mp4_path}: {e}")
            if os.path.exists(tmp_txt_path):
                os.remove(tmp_txt_path)

if __name__ == '__main__':
    main() 